#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
from typing import List

from agentscope.agent import AgentBase
from agentscope.message import Msg
from agentscope.model import DashScopeChatModel
from agentscope.formatter import DashScopeChatFormatter
from agentscope.memory import InMemoryMemory

from src.utils.env_config import aliyun_model_config

class OutlineBuilderAgent(AgentBase):
    """T4 - 报告框架搭建智能体"""

    def __init__(self):
        super().__init__()
        self.name = "OutlineBuilderAgent"
        self.sys_prompt = """你是一个报告框架构建智能体。你的任务是根据报告主题和摘要内容，生成标准的Markdown格式报告框架。
        要求：
        1. 使用Markdown语法；
        2. 包含一级标题（# 报告主题 + "报告"）；
        3. 包含二级标题"## 摘要"，其下为提供的摘要内容；
        4. 包含二级标题"## 分点提纲"，其下列出至少3个一级提纲（如1. XXXX），每个一级提纲下包含至少2个二级子项（如1.1 XXXX）；
        5. 提纲内容需与主题高度相关，逻辑合理，体现专业性；
        6. 输出格式必须为JSON，仅包含一个字段"markdown_content"，值为字符串；
        7. 不要解释过程，不要添加额外内容。
        8. 如果文献检索结果为空或信息不足，输出的内容为'文献不足，建议补充检索'的Markdown格式报告框架
        示例输出：{
        "markdown_content": 
        "# 2025 Q3 AI 行业趋势报告\\n\\n
        ## 摘要\\n本报告基于对2025年Q3 AI行业动态的文献检索……\\n\\n
        ## 分点提纲\\n
        1. 核心技术进展\\n 1.1 大模型效率优化技术\\n 1.2 多模态融合算法创新\\n
        2. 重点应用领域\\n 2.1 智能制造中的AI质检\\n 2.2 医疗领域辅助诊断系统\\n
        3. 市场规模与竞争格局\\n 3.1 全球AI市场季度增长数据\\n 3.2 头部企业技术路线对比"
        }"""

        self.model = DashScopeChatModel(
            model_name=aliyun_model_config["model_name"],
            api_key=aliyun_model_config["api_key"],
            stream=False
        )
        self.formatter = DashScopeChatFormatter()
        self.memory = InMemoryMemory()

    async def reply(self, msg: Msg) -> Msg:
        """处理报告框架构建请求"""
        await self.memory.add(msg)

        # 准备提示
        prompt = await self.formatter.format([
            Msg("system", self.sys_prompt, "system"),
            msg
        ])

        # 调用模型
        response = await self.model(prompt)

        # 创建回复消息
        reply_msg = Msg(
            name=self.name,
            content=response.content,
            role="assistant"
        )

        await self.memory.add(reply_msg)
        return reply_msg

    async def build_outline(self, topic: str, abstract: str) -> str:
        """构建报告框架"""
        user_msg = Msg(
            name="user",
            content=f"现在开始处理：\n报告主题：{topic}\n摘要内容：{abstract}",
            role="user"
        )

        response_msg = await self.reply(user_msg)

        try:
            # 提取文本内容
            content_text = ""
            if isinstance(response_msg.content, list):
                for block in response_msg.content:
                    if hasattr(block, 'text'):
                        content_text += block.text
                    elif isinstance(block, dict) and 'text' in block:
                        content_text += block['text']
            else:
                content_text = str(response_msg.content)

            result = json.loads(content_text)
            return result.get("markdown_content", "")
        except json.JSONDecodeError:
            # 如果JSON解析失败，返回默认框架
            print(f"JSON解析失败:{response_msg.content}")
            return f"""文献不足，建议补充检索"""

    async def build_enhanced_outline(self, topic: str, abstract: str, enhanced_content: dict) -> str:
        """构建增强型报告框架（整合T5/T6输出）"""

        # 构建增强的系统提示词
        enhanced_sys_prompt = """你是一个增强型报告框架构建智能体。你的任务是根据报告主题、摘要内容以及可选的技术案例分析和政策影响评估，生成完整的Markdown格式报告框架。
        要求：
        1. 使用Markdown语法；
        2. 包含一级标题（# 报告主题 + "报告"）；
        3. 包含二级标题"## 摘要"，其下为提供的摘要内容；
        4. 如果提供了技术案例分析，添加"## 技术案例分析"章节；
        5. 如果提供了政策影响评估，添加"## 政策影响评估"章节；
        6. 包含二级标题"## 分点提纲"，其下列出至少3个一级提纲，每个一级提纲下包含至少2个二级子项；
        7. 提纲内容需与主题高度相关，并结合技术案例和政策分析的内容；
        8. 输出格式必须为JSON，仅包含一个字段"markdown_content"，值为字符串；
        9. 不要解释过程，不要添加额外内容。
        示例输出：
        {
        "markdown_content": "# 2025 Q3 AI 行业趋势报告\\
        ## 摘要\
        本报告基于对2025年Q3 AI行业动态的文献检索……\\
        ## 技术案例分析\
        [技术案例内容]\\## 政策影响评估\
        [政策影响内容]\\
        ## 分点提纲\
        1. 核心技术进展\
        1.1 大模型效率优化技术\
        1.2 多模态融合算法创新\\
        2. 重点应用领域\
        2.1 智能制造中的AI质检\
        2.2 医疗领域辅助诊断系统\\
        3. 市场规模与竞争格局\
        3.1 全球AI市场季度增长数据\
        3.2 头部企业技术路线对比"
        }"""

        # 构建用户输入
        user_content = f"现在开始处理：报告主题：{topic}摘要内容：{abstract}"
        
        if "case_analysis" in enhanced_content:
            user_content += f"技术案例分析：{enhanced_content['case_analysis']}"
        
        if "policy_impact" in enhanced_content:
            user_content += f"政策影响评估：{enhanced_content['policy_impact']}"

        # 临时更新系统提示词
        original_sys_prompt = self.sys_prompt
        self.sys_prompt = enhanced_sys_prompt
        
        try:
            user_msg = Msg(
                name="user",
                content=user_content,
                role="user"
            )

            response_msg = await self.reply(user_msg)

            # 提取文本内容
            content_text = ""
            if isinstance(response_msg.content, list):
                for block in response_msg.content:
                    if hasattr(block, 'text'):
                        content_text += block.text
                    elif isinstance(block, dict) and 'text' in block:
                        content_text += block['text']
            else:
                content_text = str(response_msg.content)

            result = json.loads(content_text)
            return result.get("markdown_content", "")
            
        except json.JSONDecodeError:
            # 如果JSON解析失败，返回默认增强框架
            print(f"增强框架JSON解析失败:{response_msg.content}")
            return ""

    async def observe(self, msg: Msg) -> None:
        """观察消息"""
        await self.memory.add(msg)

    async def handle_interrupt(self) -> Msg:
        """处理中断"""
        return Msg(
            name=self.name,
            content="报告框架搭建任务被中断",
            role="assistant"
        )
